#!/usr/bin/python
# -*- coding: utf-8 -*-
# Description: 文件辅助方法
# Create: 2008-6-25
# Author: MK2[fengmk2@gmail.com]
import os
import struct
from const import FileType

# Archive
ZIP_HEADER_SIGNATURE = 0x04034b50 #zip
RAR_HEADER_SIGNATURE = 0x21726152 #rar
CAB_HEADER_SIGNATURE = 0x4643534D #cab
# Image
JPG_HEADER_SIGNATURE = 0xD8FF     #jpg
BMP_HEADER_SIGNATURE = 0x4D42 #bmp
GIF8_HEADER_SIGNATURE = 0x38464947 #gif8x
PNG_HEADER_SIGNATURE = 0x0A1A0A0D474E5089 #png
# file
PDF_HEADER_SIGNATURE = 0x46445025 #%PDF

header_signature_format = (
       (FileType.ZIP, ZIP_HEADER_SIGNATURE, 'L', 4),
       (FileType.RAR, RAR_HEADER_SIGNATURE, 'L', 4),
       (FileType.JPG, JPG_HEADER_SIGNATURE, 'H', 2),
       (FileType.BMP, BMP_HEADER_SIGNATURE, 'H', 2),
       (FileType.GIF, GIF8_HEADER_SIGNATURE, 'L', 4),
       (FileType.PNG, PNG_HEADER_SIGNATURE, 'Q', 8),
       (FileType.CAB, CAB_HEADER_SIGNATURE, 'L', 4),
       (FileType.PDF, PDF_HEADER_SIGNATURE, 'L', 4),
        )

def get_type(datas):
    for filetype, signature, format, size in header_signature_format:
        try:
            header_signature = struct.unpack(format, datas[:size])[0]
            if header_signature == signature:
                return filetype
        except struct.error:
            pass
    return FileType.Unknow
    
def get_filetype(filepath):
    """根据各类型文件格式获取文件类型, PS:只是估计，无法100%确定，但能100%否定。
    return filetype
    
    >>> path = 'D:/testsamples'
    >>> FileType.Nofile == get_filetype(path + '/empty/awfafafwafawfw.dll')
    True
    >>> FileType.Empty == get_filetype(path + '/empty/ee017db73ed184cc0aabd2dad32be81e.dll')
    True
    >>> FileType.Unknow == get_filetype(path + '/unknow/f048d0f3907d69f54ce405ce8bff3016.dll')
    True
    >>> FileType.RAR == get_filetype(path + '/rar/zeng7.rar')
    True
    >>> FileType.ZIP == get_filetype(path + '/zip/curl-7.18.2-win32-ssl-sspi.zip')
    True
    >>> FileType.CAB == get_filetype(path + '/cab/engine32.cab')
    True
    >>> FileType.JPG == get_filetype(path + '/jpg/3c9f8d47a25240252cdfc11192b2da5c.jpg')
    True
    >>> FileType.BMP == get_filetype(path + '/bmp/test.bmp')
    True
    >>> FileType.GIF == get_filetype(path + '/gif/sdfsfd.gif')
    True
    >>> FileType.PNG == get_filetype(path + '/png/a2_005.png')
    True
    >>> FileType.Empty == get_filetype(path + '/empty/blank.txt')
    True
    >>> FileType.Unknow == get_filetype(path + '/unknow/2byte.txt')
    True
    >>> FileType.PDF == get_filetype(path + '/pdf/asn1berder.pdf')
    True
    """
    if not os.path.exists(filepath) or not os.path.isfile(filepath):
        return FileType.Nofile
    if os.path.getsize(filepath) == 0:
        return FileType.Empty
#    result, filetype = peutil.is_pe(filepath)
    try:
        f = open(filepath, 'rb')
        try:
            for filetype, signature, format, size in header_signature_format:
                f.seek(0)
                try:
                    header_signature = struct.unpack(format, f.read(size))[0]
                    if header_signature == signature:
                        return filetype
                except struct.error:
                    pass
            return FileType.Unknow
        finally:
            f.close()
    except IOError, e:
        raise 

if os.name == 'nt':
    FolderExistsError = (WindowsError, OSError)
    FolderExistsCode = (183, 17)
elif os.name == 'posix':
    FolderExistsError = (OSError,)
    FolderExistsCode = (17,)
    
def ensure_dir_exist(dirpath):
    if not os.path.exists(dirpath):
        try:
            os.makedirs(dirpath)
        except FolderExistsError, e:
            if not e.args[0] in FolderExistsCode:
                raise

def __doc_test():
    import doctest
    doctest.testmod()
    
if __name__ == '__main__':
    __doc_test()